from rest_framework.exceptions import APIException
from rest_framework.response import Response
from rest_framework.views import exception_handler


def custom_exception_handler(exc, context):
    # 判断是否为APIException的子类，如果是，则进行处理
    if isinstance(exc, APIException):
        # 将异常转化为APIResponse对象
        response = APIResponse(code=exc.status_code, msg=exc.detail, data=None)
        # 设置状态码
        response.status_code = exc.status_code
        return response

    # 如果不是，则交给DRF的默认异常处理方式处理
    return exception_handler(exc, context)


class APIResponse(Response):
    def __init__(self, code=0, msg='成功', data=None, status=None,
                 template_name=None, headers=None, exception=False, content_type=None):
        """
        :param code: 0 成功 -1 操作失败
        :param msg:
        :param data:
        :param status:
        :param template_name:
        :param headers:
        :param exception:
        :param content_type:
        """

        response_data = {
            'code': code,
            'msg': msg,
            'data': data
        }
        super().__init__(response_data, status=status, template_name=template_name,
                         headers=headers, exception=exception, content_type=content_type)
